package canary

import (
	"net/http"
	"strings"
	"sync"
)

type Rule struct {
	Header    string            `yaml:"header"`
	Cookie    string            `yaml:"cookie"`
	Query     string            `yaml:"query"`
	Path      string            `yaml:"path"`
	Method    string            `yaml:"method"`
	Weight    int               `yaml:"weight"`
	Backend   string            `yaml:"backend"`
	Headers   map[string]string `yaml:"headers"`
	URLParams map[string]string `yaml:"url_params"`
}

type CanaryMatcher struct {
	rules []Rule
	mu    sync.RWMutex
}

func NewCanaryMatcher(rules []Rule) *CanaryMatcher {
	return &CanaryMatcher{
		rules: rules,
	}
}

func (cm *CanaryMatcher) UpdateRules(rules []Rule) {
	cm.mu.Lock()
	defer cm.mu.Unlock()
	cm.rules = rules
}

func (cm *CanaryMatcher) Match(r *http.Request) *Rule {
	cm.mu.RLock()
	defer cm.mu.RUnlock()

	for _, rule := range cm.rules {
		if cm.matchRule(r, &rule) {
			return &rule
		}
	}
	return nil
}

func (cm *CanaryMatcher) matchRule(r *http.Request, rule *Rule) bool {
	// 匹配 Header
	if rule.Header != "" {
		parts := strings.SplitN(rule.Header, ":", 2)
		if len(parts) != 2 {
			return false
		}
		if r.Header.Get(strings.TrimSpace(parts[0])) != strings.TrimSpace(parts[1]) {
			return false
		}
	}

	// 匹配 Cookie
	if rule.Cookie != "" {
		cookie, err := r.Cookie(rule.Cookie)
		if err != nil || cookie == nil {
			return false
		}
	}

	// 匹配 Query 参数
	if rule.Query != "" {
		parts := strings.SplitN(rule.Query, "=", 2)
		if len(parts) != 2 {
			return false
		}
		if r.URL.Query().Get(parts[0]) != parts[1] {
			return false
		}
	}

	// 匹配 Path
	if rule.Path != "" && !strings.HasPrefix(r.URL.Path, rule.Path) {
		return false
	}

	// 匹配 Method
	if rule.Method != "" && r.Method != rule.Method {
		return false
	}

	return true
}

type CanaryModifier interface {
	ModifyRequest(r *http.Request, rule *Rule) error
}

type HeaderModifier struct{}

func (h *HeaderModifier) ModifyRequest(r *http.Request, rule *Rule) error {
	for k, v := range rule.Headers {
		r.Header.Set(k, v)
	}
	return nil
}

type URLParamModifier struct{}

func (u *URLParamModifier) ModifyRequest(r *http.Request, rule *Rule) error {
	q := r.URL.Query()
	for k, v := range rule.URLParams {
		q.Set(k, v)
	}
	r.URL.RawQuery = q.Encode()
	return nil
}

type CanaryMiddleware struct {
	matcher     *CanaryMatcher
	modifiers   []CanaryModifier
	nextBackend string
}

func NewCanaryMiddleware(rules []Rule, nextBackend string) *CanaryMiddleware {
	return &CanaryMiddleware{
		matcher: NewCanaryMatcher(rules),
		modifiers: []CanaryModifier{
			&HeaderModifier{},
			&URLParamModifier{},
		},
		nextBackend: nextBackend,
	}
}

func (cm *CanaryMiddleware) Handle(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		rule := cm.matcher.Match(r)
		if rule != nil {
			// 应用规则修改
			for _, modifier := range cm.modifiers {
				if err := modifier.ModifyRequest(r, rule); err != nil {
					http.Error(w, "Internal Server Error", http.StatusInternalServerError)
					return
				}
			}
		}
		next.ServeHTTP(w, r)
	})
}
